package demo.db.index;

/**
 * B树
 */
public class BTree {
    BTreeNode root;
    int t;  // 最小度数

    public BTree(int t) {
        this.t = t;
        root = null;
    }

    public BTreeNode search(int key) {
        return (root == null) ? null : root.search(key);
    }

    public void insert(int key) {
        if (root == null) {
            root = new BTreeNode(t, true);
            root.keys[0] = key;
            root.keyCount = 1;
        } else {
            if (root.keyCount == 2 * t - 1) {
                // 根节点已满，需要分裂
                BTreeNode s = new BTreeNode(t, false);
                s.children[0] = root;
                s.splitChild(0, root);

                // 决定哪个子节点将接收新关键字
                int i = 0;
                if (s.keys[0] < key) {
                    i++;
                }
                s.children[i].insertNonFull(key);

                root = s;
            } else {
                root.insertNonFull(key);
            }
        }
    }
}

class BTreeNode {
    int[] keys;          // 关键字数组
    int t;               // 最小度数(定义节点中关键字数量的范围)
    BTreeNode[] children;// 子节点数组
    int keyCount;        // 当前关键字数量
    boolean isLeaf;      // 是否为叶子节点

    // 构造函数
    public BTreeNode(int t, boolean isLeaf) {
        this.t = t;
        this.isLeaf = isLeaf;
        keys = new int[2 * t - 1];    // 最多2t-1个关键字
        children = new BTreeNode[2 * t]; // 最多2t个子节点
        keyCount = 0;
    }

    // 查找关键字
    public BTreeNode search(int key) {
        int i = 0;
        while (i < keyCount && key > keys[i]) {
            i++;
        }

        if (i < keyCount && keys[i] == key) {
            return this;  // 在当前节点找到
        }

        if (isLeaf) {
            return null;    // 未找到
        }

        return children[i].search(key); // 递归查找子节点
    }

    // 插入关键字到非满节点
    public void insertNonFull(int key) {
        int i = keyCount - 1;

        if (isLeaf) {
            // 叶子节点：找到位置并插入
            while (i >= 0 && keys[i] > key) {
                keys[i + 1] = keys[i];
                i--;
            }
            keys[i + 1] = key;
            keyCount++;
        } else {
            // 内部节点：找到合适的子节点
            while (i >= 0 && keys[i] > key) {
                i--;
            }

            if (children[i + 1].keyCount == 2 * t - 1) {
                // 子节点已满，需要分裂
                splitChild(i + 1, children[i + 1]);
                if (keys[i + 1] < key) {
                    i++;
                }
            }
            children[i + 1].insertNonFull(key);
        }
    }

    // 分裂子节点
    public void splitChild(int i, BTreeNode y) {
        BTreeNode z = new BTreeNode(y.t, y.isLeaf);
        z.keyCount = t - 1;

        // 复制后半部分关键字到新节点
        for (int j = 0; j < t - 1; j++) {
            z.keys[j] = y.keys[j + t];
        }

        // 如果不是叶子节点，复制子节点指针
        if (!y.isLeaf) {
            for (int j = 0; j < t; j++) {
                z.children[j] = y.children[j + t];
            }
        }

        y.keyCount = t - 1;

        // 为当前节点创建新子节点空间
        for (int j = keyCount; j >= i + 1; j--) {
            children[j + 1] = children[j];
        }
        children[i + 1] = z;

        // 移动关键字并插入中间关键字
        for (int j = keyCount - 1; j >= i; j--) {
            keys[j + 1] = keys[j];
        }
        keys[i] = y.keys[t - 1];
        keyCount++;
    }
}